package com.shxzz.easy12306.services.ticket.service.handler.tokenbucket;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.shxzz.easy12306.framework.starter.cache.StringRedisTemplateProxy;
import com.shxzz.easy12306.services.ticket.common.constant.RedisKeyConstant;
import com.shxzz.easy12306.services.ticket.dao.entity.TrainDO;
import com.shxzz.easy12306.services.ticket.dto.domain.PurchaseTicketPassengerDetailDTO;
import com.shxzz.easy12306.services.ticket.dto.domain.RouteDTO;
import com.shxzz.easy12306.services.ticket.dto.domain.SelectSeatDTO;
import com.shxzz.easy12306.services.ticket.dto.req.PurchaseTicketReqDTO;
import com.shxzz.easy12306.services.ticket.service.TrainService;
import com.shxzz.easy12306.services.ticket.service.TrainStationService;
import com.shxzz.easy12306.services.ticket.service.cache.TicketPurchaseTokenBucketCacheLoader;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * 车票购买“令牌桶”，对应车票余量
 *
 * @author shxzz
 * @create 2024-03-11-15:37
 */
@RequiredArgsConstructor
@Component
public class TicketPurchaseTokenBucket {

    private final TrainService trainService;
    private final TrainStationService trainStationService;

    private final TicketPurchaseTokenBucketCacheLoader ticketPurchaseTokenBucketCacheLoader;

    private final StringRedisTemplateProxy stringRedisTemplateProxy;

    private static final String LUA_GET_TOKEN_SCRIPT_PATH = "lua/ticket_availability_token_bucket.lua";
    private static final String LUA_ROLLBACK_TOKEN_SCRIPT_PATH = "lua/rollback_ticket_availability_token_bucket.lua";
    private DefaultRedisScript<Long> getTokenLuaScript;
    private DefaultRedisScript<Long> rollbackTokenLuaScript;

    /**
     * 初始化 Lua 脚本
     */
    @PostConstruct
    public void init() {
        DefaultRedisScript<Long> getTokenRedisScript = new DefaultRedisScript<>();
        getTokenRedisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource(LUA_GET_TOKEN_SCRIPT_PATH)));
        getTokenRedisScript.setResultType(Long.class);
        getTokenLuaScript = getTokenRedisScript;

        DefaultRedisScript<Long> rollbackTokenRedisScript = new DefaultRedisScript<>();
        rollbackTokenRedisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource(LUA_ROLLBACK_TOKEN_SCRIPT_PATH)));
        rollbackTokenRedisScript.setResultType(Long.class);
        rollbackTokenLuaScript = rollbackTokenRedisScript;
    }

    /**
     * 从令牌桶获取令牌
     *
     * @return true:获取令牌成功，false:获取令牌失败
     */
    public boolean getTokenFromBucket(PurchaseTicketReqDTO requestParam) {
        TrainDO trainInfo = trainService.getTrainInfo(requestParam.getTrainId());
        StringRedisTemplate stringRedisTemplate = stringRedisTemplateProxy.getInstance();
        String tokenBucketKey = RedisKeyConstant.PURCHASE_TICKET_TOKEN_BUCKET_PREFIX + requestParam.getTrainId();
        if (Boolean.FALSE.equals(stringRedisTemplate.hasKey(tokenBucketKey))) {
            ticketPurchaseTokenBucketCacheLoader.load(requestParam.getTrainId(), trainInfo.getTrainType(), trainInfo.getStartStation(), trainInfo.getEndStation());
        }

        Map<Integer, Long> seatTypeCountMap = requestParam.getPassengers()
                .stream()
                .collect(Collectors.groupingBy(PurchaseTicketPassengerDetailDTO::getSeatType, Collectors.counting()));
        JSONArray seatTypeCountArray = getSeatTypeCountJSONArray(seatTypeCountMap);

        String departure = requestParam.getDeparture();
        String arrival = requestParam.getArrival();
        String remainTicketPrefix = departure + "-" + arrival;

        List<RouteDTO> routeDTOList = trainStationService.listTakeoutTrainStationRoute(requestParam.getTrainId(), departure, arrival);

        Long getTokenResult = stringRedisTemplate.execute(
                getTokenLuaScript,
                List.of(tokenBucketKey, remainTicketPrefix),
                JSON.toJSONString(seatTypeCountArray),
                JSON.toJSONString(routeDTOList)
        );
        return Objects.equals(getTokenResult, 0L);
    }

    /**
     * 回滚令牌桶令牌
     */
    public void rollbackToken(List<SelectSeatDTO> selectSeatDTOList, List<RouteDTO> routeDTOList, String trainId) {
        String tokenBucketKey = RedisKeyConstant.PURCHASE_TICKET_TOKEN_BUCKET_PREFIX + trainId;
        StringRedisTemplate stringRedisTemplate = stringRedisTemplateProxy.getInstance();
        Map<Integer, Long> seatTypeCountMap = selectSeatDTOList.stream()
                .collect(Collectors.groupingBy(SelectSeatDTO::getSeatType, Collectors.counting()));
        JSONArray seatTypeCountArray = getSeatTypeCountJSONArray(seatTypeCountMap);
        stringRedisTemplate.execute(
                rollbackTokenLuaScript,
                List.of(tokenBucketKey),
                JSON.toJSONString(seatTypeCountArray),
                JSON.toJSONString(routeDTOList)
        );
    }

    private static JSONArray getSeatTypeCountJSONArray(Map<Integer, Long> seatTypeCountMap) {
        return seatTypeCountMap.entrySet().stream().map(entry -> {
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("seatType", String.valueOf(entry.getKey()));
            jsonObject.put("count", String.valueOf(entry.getValue()));
            return jsonObject;
        }).collect(Collectors.toCollection(JSONArray::new));
    }
}
